{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import base64\n",
    "from collections import deque\n",
    "import os\n",
    "import pathlib\n",
    "import shutil\n",
    "\n",
    "from IPython import display as ipydisplay\n",
    "import torch\n",
    "\n",
    "from utils_env import MyEnv\n",
    "from utils_drl import Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = 120\n",
    "model_name = f\"model_{target:03d}_best\"\n",
    "model_path = f\"./saved_models/{model_name}\"\n",
    "device = torch.device(\"cpu\")\n",
    "env = MyEnv(device)\n",
    "agent = Agent(env.get_action_dim(), device, 0.99, 0, 0, 0, 1, model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_queue = deque(maxlen=5)\n",
    "avg_reward, frames = env.evaluate(obs_queue, agent, render=True)\n",
    "print(f\"Avg. Reward: {avg_reward:.1f}\")\n",
    "\n",
    "!rm -r eval_*\n",
    "target_dir = f\"eval_{target:03d}\"\n",
    "os.mkdir(target_dir)\n",
    "for ind, frame in enumerate(frames):\n",
    "    frame.save(os.path.join(target_dir, f\"{ind:06d}.png\"), format=\"png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_to_mp4 = os.path.join(target_dir, \"movie.mp4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(path_to_mp4):\n",
    "    shutil.move(target_dir, \"tmp_eval_frames\")\n",
    "    # Generate an mp4 video from the frames\n",
    "    !ffmpeg -i \"./tmp_eval_frames/%06d.png\" -pix_fmt yuv420p -y ./tmp_eval_movie.mp4 > /dev/null 2>&1\n",
    "    !rm -r tmp_eval_frames\n",
    "    os.mkdir(target_dir)\n",
    "    shutil.move(\"tmp_eval_movie.mp4\", path_to_mp4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HTML_TEMPLATE = \"\"\"<video alt=\"{alt}\" autoplay loop controls style=\"height: 400px;\">\n",
    "  <source src=\"data:video/mp4;base64,{data}\" type=\"video/mp4\" />\n",
    "</video>\"\"\"\n",
    "\n",
    "def show_video(path_to_mp4: str) -> None:\n",
    "    \"\"\"show_video creates an HTML element to display the given mp4 video in IPython.\"\"\"\n",
    "    mp4 = pathlib.Path(path_to_mp4)\n",
    "    video_b64 = base64.b64encode(mp4.read_bytes())\n",
    "    html = HTML_TEMPLATE.format(alt=mp4, data=video_b64.decode('ascii'))\n",
    "    ipydisplay.display(ipydisplay.HTML(data=html))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "show_video(path_to_mp4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3-final"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}